'''
Created on Jul 4, 2009
@author: mjbommar
'''

from tfaces.collection.Experiment import *
from tfaces.collection.util import *


class TimeSeriesExperiment(Experiment):
    '''
    The Time Series Experiment class loads an Experiment that produces time series output, 
    such as most Drone or Gridsweeper Drone jobs.
    
    The output should be of the following form:
    
    parameter1=value
    parameter2=value
    .
    .
    .
    # field1name, field2name, field3name, ...
    data1, data2, data3, ...
    data1, data2, data3, ...
    .
    .
    .
    
    
    Any lines that begin with # that do not fall directly before the data section will be treated as comments.
    
    For more information about Drone or GridSweeper, see the following links:
        Drone: http://www.cscs.umich.edu/Software/Drone/
        GridSweeper: http://cscs.umich.edu/PmWiki/pmwiki.php/CSCSSoftware/GridSweeper
    '''
    
    droneParameters = ['reportFileName', 'reportFrequency', 'outputDirName', 'runNumber', 'initialParametersFileName', 'initialAgentsFileName', 'saveRunEndState', 'reportFileNameSuffix']
    
    def __init__(self, experimentDirectory='.', reportFilePattern='report.[0-9]+'):
        '''
        The Drone Experiment constructor takes as arguments the directory in which the
        experiment was stored, and whether to parse the headers from the report file.
        '''

        # Call parent constructor
        Experiment.__init__(self, experimentDirectory, reportFilePattern)
        
        self.processTextRuns(self.findRuns())
        
    def processTextRuns(self, runFiles):
        """
        This method processes a list of text run files.
        """
        runs = [self.processTextRun(runFile) for runFile in runFiles]
        self.addRuns(runs)

    def processTextRun(self, runFile):
        """
        Parse each report file and return a Run object.
        """
        # Initialize the variables for the Run object
        parameters = {}
        outputs = {}
        seed = 0

        # Load the data and split it into lines
        data = open(runFile).read()
        dataLines = [line.strip() for line in data.strip().splitlines()]
        seenHeader = False
        dataRows = []

        for lineNo, line in enumerate(dataLines):
            # Skip header lines
            if line.startswith('#') or line.startswith('@'):
                continue
            
            # Look for parameter definitions
            if '=' in line:
                name, value = [token.strip() for token in line.split('=')]
                if not isOneOf(value, [str, int, float]):
                    raise Exception("TimeSeriesExperiment::processTextRun: invalid definition type for %s on line %d" % (name, lineNo))
                
                if name == 'seed':
                    seed = int(value)
                elif name not in self.droneParameters:
                    parameters[name] = toNaturalType(value)

            elif len(line) > 0:
                # Check to see if we've passed what should be the header line
                if not seenHeader:
                    # Parse the header from the previous line
                    headerTokens = self.lineSplit(dataLines[lineNo - 1])[1:]
                    seenHeader = True
                
                # Now parse the data for each line and make sure it matches the header
                dataTokens = self.lineSplit(line)
                if len(dataTokens) != len(headerTokens):
                    raise Exception("TimeSeriesExperiment::processTextRun: number of data columns not equal to number of headers on line %d" % lineNo)
                dataRows.append(dataTokens)
        
        # Produce the actual output dictionary by forcing the data from the file into its natural type
        for i, header in enumerate(headerTokens):
            outputs[header] = [toNaturalType(row[i]) for row in dataRows]

        return TimeSeriesRun(parameters, outputs, seed, False)
    
    def initDB(self):
        '''
        Initialize the database
        '''
        
        # Generate the parameter set table from the parameter list.
        sqlParameterSet = 'CREATE TABLE IF NOT EXISTS parameterset ('
        for parameter in self.parameterKeys:
            sqlParameterSet += '%s TEXT,' % parameter
        sqlParameterSet += 'id INTEGER PRIMARY KEY AUTOINCREMENT)'
        
        self.db.execute(sqlParameterSet)
        
        # Generate the run table from the output list.
        sqlRuns = 'CREATE TABLE IF NOT EXISTS runs ('
        for output in self.outputKeys:
            sqlRuns += '%s REAL,' % output
        sqlRuns += 'seed INTEGER NOT NULL, parametersetid INTEGER NOT NULL)'
        
        self.db.execute(sqlRuns)
        self.db.commit()
    
    def getParameterSetID(self, parameters):
        '''
        Determine the parameter set ID from the given dictionary.
        '''
        
        # Initialize the SELECT and INSERT strings
        selectString = "SELECT id FROM parameterset WHERE "
        insertString = "INSERT INTO parameterset VALUES ("
        
        # Generate the SELECT and INSERT strings from the parameters
        for i, (k, v) in enumerate(parameters.iteritems()):
            if i < len(parameters) - 1:
                selectString += '%s = "%s" AND ' % (k, v)
                insertString += '"%s",' % v
            else:
                selectString += '%s = "%s"' % (k, v)
                insertString += '"%s", NULL)' % v
        
        ''' Run the SELECT string, see if the row exists.  If it exists, return the ID.
        Otherwise, INSERT the row, then return the new ID.'''
        result = self.db.execute(selectString).fetchone()
        if result:
            return result[0]
        else:
            self.db.execute(insertString)
            return self.db.execute(selectString).fetchone()[0]
    
    def save(self, fileName):
        '''
        Write the experiment out to a SQLite database file.
        '''
        # If the experiment file already exists, raise an exception.
        if os.path.exists(fileName):
            raise Exception("Experiment::save: %s already exists" % (fileName))
        
        # Initialize the database.
        self.db = sqlite3.connect(fileName)
        self.initDB()
        
        
        sqlInsertData = []
        
        # Generate the SQL queries for the data for all runs.  
        for run in self.runs:
            ''' Get the parameter set ID for the given run.  Then, we
            generate the template for the INSERT statement.''' 
            parameterSetID = self.getParameterSetID(run.parameters)
            baseInsertString = 'INSERT INTO runs ('
            for i, output in enumerate(run.outputs.keys()):
                if i < len(run.outputs.keys()) - 1:
                    baseInsertString += '%s,' % output
                else:
                    baseInsertString += '%s, seed, parametersetid) VALUES (' % output
            
            for i, output in enumerate(run.outputs.keys()):
                if i < len(run.outputs.keys()) - 1:
                    baseInsertString += '?,'
                else:
                    baseInsertString += '?,?,?)'
            
            ''' Now generate the data tuples that we pass to the template INSERT
            statement in batch mode. '''
            for i in range(run.numRows):
                sqlData = []
                for j, output in enumerate(run.outputs.keys()):
                    if j < len(run.outputs.keys()) - 1:
                        sqlData.append(run.outputs[output][i])
                    else:
                        sqlData.extend([run.outputs[output][i], run.seed, parameterSetID])
                sqlInsertData.append(tuple(sqlData))
        
        # Finally add the data and commit.
        self.db.executemany(baseInsertString, sqlInsertData)
        self.numrows = len(sqlInsertData)
        self.db.commit()
        self.db.close()
    
